package com.hsogoo.collate.datasource;

import com.google.common.collect.Maps;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

/**
 * @Author huangsaigang
 * Created @2019-10-25.
 */
@Slf4j
public class DynamicDataSource extends AbstractRoutingDataSource {

    /**
     * 获取目标数据源
     */
    @Override
    protected DataSource determineTargetDataSource() {

        //如果不存在则表示，只在DataSourceHolder里添加的数据源，并未加入key
        if (!DynamicDataSourceContextHolder.containDataSourceKey(DynamicDataSourceContextHolder.getDataSourceKey())) {
            //初始化外部数据源
            Map<String, DataSource> dataSources = DataSourceHolder.getDataSources();
            //加载到动态数据源的targetDataSources
            setDataSources(dataSources);
        }
        DataSource targetDataSource = DataSourceHolder.getDataSource(DynamicDataSourceContextHolder.getDataSourceKey());
        if (targetDataSource == null) {
            String exception = String.format("DataSource [%s] doesn't exist", DynamicDataSourceContextHolder.getDataSourceKey());
            throw new RuntimeException(exception) ;
        }
        log.info("DynamicDataSource current thread id {}", Thread.currentThread().getId());
        return targetDataSource;
    }

    /**
     * 如何路由数据源
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return DynamicDataSourceContextHolder.getDataSourceKey();
    }

    /**
     * 设置默认数据源
     * @param defaultDataSource
     */
    public void setDefaultDataSource(Object defaultDataSource) {
        super.setDefaultTargetDataSource(defaultDataSource);
        DataSourceHolder.addDataSource(DynamicDataSourceContextHolder.DEFAULT_DATASOURCE, (DataSource)defaultDataSource);
    }

    /**
     * 设置数据源
     * @param dataSources
     */
    public void setDataSources(Map<String, DataSource> dataSources) {
        setSuperTargetDataSources(dataSources);
        // 将数据源的 key 放到数据源上下文的 key 集合中，用于切换时判断数据源是否有效
        DynamicDataSourceContextHolder.addDataSourceKeys(dataSources.keySet());
    }

    private void setSuperTargetDataSources(Map<String, DataSource> dataSources){
        Map<Object, Object> targetDataSources = Maps.newHashMap();
        for(Map.Entry<String, DataSource> entry : dataSources.entrySet()){
            targetDataSources.put(entry.getKey(), entry.getValue());
        }
        super.setTargetDataSources(targetDataSources);
    }

}
